#include <cmath>
#include <iostream>
#include <stdlib.h>
#include <stdio.h>
#include <string>
#include <mpi.h>
#include <unistd.h>
#define A1 -1
#define A2 2
#define B1 -2
#define B2 2
#define EPS 0.000001
#define DOWN_TAG 100
#define DEBUG 0

using namespace std;

void get_cmd_params(int _argc, char **_argv, int &_M, int &_N)
{
    // Параметры размерности по умолчанию
    _M = 20;
    _N = 20;
    
    // Обработка параметров командной строки
    for (int i = 1; i < _argc; i++)
    {
        string option(_argv[i]);
        
        if (option.compare("-M") == 0)
        {
            _M = atoi(_argv[++i]);
        }

        if (option.compare("-N") == 0)
        {
            _N = atoi(_argv[++i]);
        }
    }

}

double func_u(double x, double y){
    return exp(1-(x+y)*(x+y));
}

double func_q(double x, double y){
    if (x+y < 0)
        return 0;
    else
        return (x+y)*(x+y);
}

double func_k(double x, double y){
    return 4+x;
}

double func_F(double x, double y){
    return 2*(8 + 3*x + y - 4*(4+x)*(x+y)*(x+y))*func_u(x, y) + func_q(x, y)*func_u(x, y);
}
 
void vector_diff(double* res, double *w1, double *w0, int M, int N){       // Разность двух векторов
    int i, j;

    //# pragma omp parallel for private(i, j)
    for(i = 0; i < M+2; i++)
        for (j = 0; j < N+2; j++)
                if(i == 0 or i == (M+1) or j == 0 or j == (N+1))
                    res[i*(N+2) + j] = 0;
                else
                    res[i*(N+2) + j] = w1[i*(N+2) + j] - w0[i*(N+2) + j];

}

 double ro_i(int i, int M, int N){                           // Вспомогательные функции для скалярного произведения в пространстве H
    if (i == 1 or i == M+1)
        return 0.25;                                       // Попробуем поставить везде параметр 0.25 вместо 0.5, это существенно уменьшает число итераций
    else
        return 1;
}

double ro_j(int j, int M, int N){
    if (j == 1 or j == N+1)
        return 0.25;
    else
        return 1;
}

double vector_mult(double *w0, double *w1, int M, int N, double hx, double hy){          // Скалярное произведение в пространстве H: dim(H) =[ 1..(M-1) 1..(N-1) ]
    double res = 0;
    int i, j;

    //# pragma omp parallel for private(i, j)
    for(i = 0; i < M+2; i++)
        for (j = 0; j < N+2; j++)
            if(i == 0 or i == (M+1) or j == 0 or j == (N+1))
                res += 0;
            else
                res += ro_i(i, M, N)*ro_j(j, M, N)*w0[i*(N+2) + j]*w1[i*(N+2) + j];

    return res*hx*hy;
}

double norm(double *w, int M, int N, double hx, double hy){                              // Норма вектора без учёта границ
    return sqrt(vector_mult(w, w, M, N, hx, hy));
}

void B_init(double* B, int M, int N, double hx, double hy, double x_start, double y_start){                                      // Правая часть уравнения F_ij
    int i, j;

    # pragma omp parallel for private(i, j)
    for(i = 0; i < M+2; i++)
        for (j = 0; j < N+2; j++){ 
            B[i*(N+2) + j] = func_F(x_start + i*hx, y_start + j*hy);
        }

}

void A_mult(double* A, double *w, int M, int N, double hx, double hy, double x_start, double y_start){     // Хитрая функция, результатом которой будет умножение вектора на матрицу А слева
    double wx, wy;
    int i, j;

    for(i = 0; i < M+2; i++)
        for (j = 0; j < N+2; j++){
            if(i == 0 or i == (M+1) or j == 0 or j == (N+1)){  
                A[i*(N+2) + j] = w[i*(N+2) + j];         // Граничные условия сохраняем неизменными
            }
            else{
                wx = func_k(x_start + (i+0.5)*hx, y_start + j*hy)*(w[(i+1)*(N+2) + j] - w[i*(N+2) + j])/hx - func_k(x_start + (i-0.5)*hx, y_start + j*hy)*(w[i*(N+2) + j] - w[(i-1)*(N+2) + j])/hx;
                wy = func_k(x_start + i*hx, y_start + (j+0.5)*hy)*(w[i*(N+2) + j+1] - w[i*(N+2) + j])/hy - func_k(x_start + i*hx, y_start + (j-0.5)*hy)*(w[i*(N+2) + j] - w[i*(N+2) + j-1])/hy;
                A[i*(N+2) + j] = -wx/hx - wy/hy + func_q(x_start + i*hx, y_start + j*hy) * w[i*(N+2) + j];
            }
    }
                
}
    

int main(int argc, char **argv)
{   
    int M, N;

    get_cmd_params(argc, argv, M, N);

    int rank;         // ID процесса 
    int commSize;     // Общее число процессов 
    int rank_n;
    int dims[2];
    dims[0] = 0;
    dims[1] = 0;
    int write[1];
    write[0] = 0;

    // Шаги сетки
    double hx = double(A2-A1) / M;
    double hy = double(B2-B1) / N;
    double eps_global = 1.0;

    
    // Инициализация процессов, можно передать argc и argv, тут каждый сам начинает все делать
    MPI_Init(&argc, &argv);
    MPI_Status status;
    MPI_Request request;
    MPI_Comm MPI_COMM_CART;
    // Получаем ранг процесса текущего
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    // Получаем общее число процессов в системе
    MPI_Comm_size(MPI_COMM_WORLD, &commSize);
    
    // Разбиваем число процессоров по сетке размерности два и 
    // записываем число процессов по кажому направлению в dims
    MPI_Dims_create(commSize, 2, dims);

    if (DEBUG== 1) cout << rank << " : " << " DIMS " << dims[0] << " " << dims[1] << "\n";

    // Создадим коммуникатор с картезианской топологией 
    int periods[2];
    periods[0] = 0;
    periods[1] = 0;  
    MPI_Cart_create(MPI_COMM_WORLD, 2, dims, periods, true, &MPI_COMM_CART);
    // Для индексов решетки процессов
    int coords[2];
    int coords_n[2];
    MPI_Cart_coords(MPI_COMM_CART, rank, 2, coords);

    if (DEBUG == 1) cout << rank << " : " << " Coords " << coords[0] << " " << coords[1] << "\n";

    // Для каждого процесса задаем начальные данные
    // Сдвиг и размер по x (i)
    int x_shift;
    int x_size;
    int y_shift;
    int y_size;

    // Делим число процессоров по внутренней сетке расчётов
    if ((M) % dims[0] == 0){
        x_size = (M) / dims[0];
        x_shift = coords[0] * ((M) / dims[0]);
    } else {
        if (coords[0] == 0){
            x_size = (M) % dims[0] + (M) / dims[0];
            x_shift = 0;
        } else {
            x_size = (M) / dims[0];
            x_shift = (M) % dims[0] + (coords[0]) * ((M) / dims[0]);
        }
    }

    if ((N) % dims[1] == 0){
        y_size = (N) / dims[1];
        y_shift = coords[1] * ((N) / dims[1]);
    } else {
        if (coords[1] == 0){
            y_size = (N) % dims[1] + (N) / dims[1];
            y_shift = 0;
        } else {
            y_size = (N) / dims[1];
            y_shift = (N) % dims[1] + (coords[1]) * ((N) / dims[1]);
        }
    }
    if (DEBUG == 1) cout << rank << " x_shift " << x_shift << " y_shift " << y_shift << " x_size " << x_size << " y_size " << y_size << "\n"; 


    if (DEBUG == 1) cout << rank << " Start program\n";
    
    double start_time = MPI_Wtime();

    // Выделяем память для пересылки верхней границ (send и receive) размерность
    double *top_s, *top_r;
    top_s = new double[x_size];
    top_r = new double[x_size];
    // Выделяем память для пересылки нижней границы
    double *bottom_s, *bottom_r;
    bottom_s = new double[x_size];
    bottom_r = new double[x_size];
    // Выделяем память для пересылки левой границы размерность N
    double *left_s, *left_r;
    left_s = new double[y_size];
    left_r = new double[y_size];
    // Выделяем память для пересылки правой границы
    double *right_s, *right_r;
    right_s = new double[y_size];
    right_r = new double[y_size];

    int t,i,j;
    int k = 0;
           
    // Отсылаем результаты, потом получаем, делаем ввычисление

    // Выделям общий рабочий объем
    if (DEBUG == 1) cout << rank << " Start alloc memory\n";

    // w -- вытянутая в вектор матрица по столбцам
    double *w = new double [(x_size+2)*(y_size+2)];
    double *w_pr = new double [(x_size+2)*(y_size+2)];
    double *B = new double [(x_size+2)*(y_size+2)];
    double tau; 
    double eps_local, eps_r;

    double *Aw = new double [(x_size+2)*(y_size+2)];
    double *r = new double [(x_size+2)*(y_size+2)];
    double *Ar = new double [(x_size+2)*(y_size+2)];
    double *w_w_pr = new double [(x_size+2)*(y_size+2)];

    if (DEBUG == 1) cout << rank << " Start initiallizing\n";
    // Задаем начальные значения 
    # pragma omp parallel for private(i, j)                 // На первой итерации значение не имеет значения
    for (i = 0; i < x_size + 2; i++)
        for (j = 0; j < y_size + 2; j++)                
            w[i*(y_size+2) + j] = 0;

    if (DEBUG == 1) cout << rank << " Start solving\n";

   // pragma внутри функции B_init
    B_init(B, x_size, y_size, hx, hy, A1 + x_shift*hx, B1 + y_shift*hy);

    // Численный метод решения на шаге 0 norm_w = 1.0
 
    while (eps_global > EPS){
        k++;

        //# pragma omp parallel for private(i, j)
        for (i = 1; i < x_size + 1; i++){           
            for (j = 1; j < y_size + 1; j++)
                w_pr[i*(y_size+2) + j] = w[i*(y_size+2) + j];                                 
        }               

        int c;
        int t = 0; // зачем t

        // Отдаем верхнюю границу, если есть кому
        if (dims[1] != 1){
            c = 0;
            for (i = 1; i < x_size + 1; i++){
                top_s[c] = w[i*(y_size+2) + y_size]; // Предпоследнюю строчку 
                c++;
            }
            coords_n[0] = coords[0];             // neighbor's coord
            if (coords[1] != dims[1] - 1){       
                coords_n[1] = coords[1] + 1;
                
                MPI_Cart_rank(MPI_COMM_CART, coords_n, &rank_n);
                MPI_Isend(top_s, x_size, MPI_DOUBLE, rank_n, t, MPI_COMM_CART, &request);
            }
        }

        // Отдаем нижнюю границу, если есть кому
        if (dims[1] != 1){
            c = 0;
            for (i = 1; i < x_size + 1; i++){
                    bottom_s[c] = w[i*(y_size+2) + 1];
                    c++;
            }
            coords_n[0] = coords[0];
            if (coords[1] != 0){ 
                coords_n[1] = coords[1] - 1;
             
                MPI_Cart_rank(MPI_COMM_CART, coords_n, &rank_n);
                MPI_Isend(bottom_s, x_size, MPI_DOUBLE, rank_n, t + DOWN_TAG, MPI_COMM_CART, &request);

            }         
        }

        // Отдаем левую границу, если есть кому
        if (dims[0] != 1){
            if (coords[0] != 0){ 
                c = 0;
                for (j = 1; j < y_size + 1; j++){
                        left_s[c] = w[(x_size + 2)*(y_size+2) - 2*(y_size + 2) + j];
                        c++;
                }
                coords_n[0] = coords[0] - 1;
                coords_n[1] = coords[1];
                MPI_Cart_rank(MPI_COMM_CART, coords_n, &rank_n);
                MPI_Isend(left_s, y_size, MPI_DOUBLE, rank_n, t, MPI_COMM_CART, &request);
            }
        }

        // Отдаем правую границу, если есть кому
        if (dims[0] != 1){
            if (coords[0] != dims[0] - 1){    
                c = 0;
                for (j = 1; j < y_size + 1; j++){
                        right_s[c] = w[(y_size+2) + j];
                        c++;
                }
                coords_n[0] = coords[0] + 1;
                coords_n[1] = coords[1];
                MPI_Cart_rank(MPI_COMM_CART, coords_n, &rank_n);
                MPI_Isend(right_s, y_size, MPI_DOUBLE, rank_n, t, MPI_COMM_CART, &request);
            }
        }

        // Принимаем верхнюю границу, если есть от кого
        if (dims[1] != 1){
            coords_n[0] = coords[0];
            if (coords[1] != dims[1] - 1){
                coords_n[1] = coords[1] + 1;
             
                MPI_Cart_rank(MPI_COMM_CART, coords_n, &rank_n);
                // Адрес откуда, размер, тип, кому, тег, пространство
                MPI_Recv(top_r, x_size, MPI_DOUBLE, rank_n, t+DOWN_TAG, MPI_COMM_CART, &status);   
                c = 0;
                for (i = 1; i < x_size + 1; i++){
                    w[i*(y_size+2) + 0] = top_r[c];
                    c++;
                    }
           }
        } else {
            // Заполняем нижнюю границу
            for (i = 1; i < x_size + 1; i++){
                w[i*(y_size+2) + 0] = func_u(A1 + (x_shift + i)*hx, B1);
            }
        }


        // Принимаем нижнюю границу, если есть от кого
        if (dims[1] != 1){
            coords_n[0] = coords[0];
            if (coords[1] != 0){  
                coords_n[1] = coords[1] - 1;
            
                MPI_Cart_rank(MPI_COMM_CART, coords_n, &rank_n);
                // Адрес откуда, размер, тип, кому, тег, пространство
                MPI_Recv(bottom_r, x_size, MPI_DOUBLE, rank_n, t, MPI_COMM_CART, &status); 
                c = 0;
                for (i = 1; i < x_size + 1; i++){
                    w[i*(y_size+2) + y_size + 1] = bottom_r[c];
                    c++;
                }
            }
        } else {
            // Заполняем верхнюю границу
            for (i = 1; i < x_size + 1; i++){
                    w[i*(y_size+2) + y_size + 1] = func_u(A1 + (x_shift + i)*hx, B2);
            }
         }

        // Принимаем левую границу, если есть от кого
        if (dims[0] != 1 && coords[0] != 0){   
            c = 0;
            coords_n[0] = coords[0] - 1;
            coords_n[1] = coords[1];
            MPI_Cart_rank(MPI_COMM_CART, coords_n, &rank_n);
            MPI_Recv(left_r, y_size, MPI_DOUBLE, rank_n, t, MPI_COMM_CART, &status); 
            for (j = 1; j < y_size + 1; j++){
                    w[j] = left_r[c];
                    c++;
            }
        } else { 
            // Заполняем правую границу
            for (j = 1; j < y_size + 1; j++){
                    w[j] = func_u(A1, B1 + (y_shift + j)*hy);
            }
        }

        // Принимаем правую границу, если есть от кого
        if (dims[0] != 1 && coords[0] != dims[0]-1){      
            c = 0;
            coords_n[0] = coords[0] + 1;
            coords_n[1] = coords[1];
            MPI_Cart_rank(MPI_COMM_CART, coords_n, &rank_n);
            MPI_Recv(right_r, y_size, MPI_DOUBLE, rank_n, t, MPI_COMM_CART, &status); 
            for (j = 1; j < y_size + 1; j++){
                    w[(x_size + 2)*(y_size+2) - (y_size + 2) + j] = right_r[c];
                    c++;
            }
        } else { 
            // Заполняем левую границу
            for (j = 1; j < y_size + 1; j++){
                    w[(x_size + 2)*(y_size+2) - (y_size + 2) + j] = func_u(A1, B1 + (y_shift + j)*hy);
            }
        }
        

        A_mult(Aw, w, x_size, y_size, hx, hy, A1 + coords[0]*hx, B1 + coords[1]*hy);

        vector_diff(r, Aw, B, x_size, y_size);

        A_mult(Ar, r, x_size, y_size, hx, hy, A1 + coords[0]*hx, B1 + coords[1]*hy);


        tau = (vector_mult(Ar, r, x_size, y_size, hx, hy))/pow(norm(Ar, x_size, y_size, hx, hy), 2);  // У каждого процеса свой итерационный параметр
        //if (DEBUG == 1) cout << rank << " tau = " << fixed << tau << " \n";

        # pragma omp parallel for private(i, j)
        for (i = 1; i < x_size + 1; i++)
            for (j = 1; j < y_size + 1; j++)
                w[i*(y_size+2) + j] = w[i*(y_size+2) + j] - tau*r[i*(y_size+2) + j]/2;

        // ||w_{k+1} - w_{k}||
        vector_diff(w_w_pr, w, w_pr, x_size, y_size);

        eps_local = norm(w_w_pr, x_size, y_size, hx, hy);

        MPI_Allreduce(&eps_local, &eps_global, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);

        //if (DEBUG == 1) cout << "eps global " << eps_global << endl;

    } 


    // Все готово, ждем всех и пишем только результат
    MPI_Barrier(MPI_COMM_WORLD);
    double end_time = MPI_Wtime();

    //cout.precision(17);

    if (rank != 0){
        MPI_Recv(write, 1, MPI_INT, rank-1, 0, MPI_COMM_WORLD, &status); 
    } else {
        // printf("TIME %f\n", end_time - start_time);
        cout << "TIME " << end_time - start_time << "\n";
    }
    // Вывести результаты 
    //cout << rank << " START WRITING\n";
    /*
    for (i = 0; i < x_size+2; i++){
        for (j = 0; j < y_size+2; j++)             
            cout << fixed << w[i*(y_size+2) + j] << " ";
        cout << endl;
    }
    */
    if (rank == 0) cout << "number of itterations: " << k << '\n';
    if (rank == 0) cout << "tau: " << tau << '\n';
    if (rank == 0) cout << "eps: " << EPS << '\n';


   // cout << rank << " FINISH WRITING\n";
    cout << flush;
    usleep(500);

    if (rank != commSize-1){
        MPI_Send(write, 1, MPI_INT, rank+1, 0, MPI_COMM_WORLD); 
    }

    //Free each sub-array
    if (DEBUG==1) cout << rank << " Free memmory\n";

    delete[] w;
    delete[] w_pr;
    delete[] B;
    delete[] r;
    delete[] Ar;
    delete[] Aw;
    delete[] w_w_pr;

    delete[] top_s;
    delete[] top_r;
    delete[] bottom_s;
    delete[] bottom_r;
    delete[] left_s;
    delete[] left_r;
    delete[] right_s;
    delete[] right_r;


    MPI_Finalize();

    return 0;
}